import pandas as pd
import pymysql
import pymysql.cursors

class MySqlUtils(object):
    def __init__(self):
        self.conn = pymysql.connect(
            host='localhost',
            user='root',
            password='sjk1234',
            db='tushare',
            charset='utf8mb4',
            cursorclass=pymysql.cursors.DictCursor
        )


class ClassIfcation(object):
    def __init__(self):
        pass

    def get_fina_indicator(self, conn):
        cursor = conn.cursor()
        query = """
            SELECT ts_code, ann_date, eps, total_revenue_ps, undist_profit_ps, gross_margin, fcff, fcfe, tangible_asset, bps, grossprofit_margin, npta FROM tushare.financial_data WHERE tushare.financial_data.ann_date BETWEEN '2023-01-01' AND '2023-12-12' 
        """
        cursor.execute(query)
        result = cursor.fetchall()
        df = pd.DataFrame(result)
        
        df = df.dropna(subset=['eps', 'total_revenue_ps', 'gross_margin', 'fcff', 'fcfe', 'tangible_asset', 'bps', 'grossprofit_margin', 'npta'])
        df = df.reset_index(drop=True)
        return df
        
    def get_daily(self, conn, df):
        cursor = conn.cursor()
        new_list = []

        for index, row in df.iterrows():
            print(index)
            ts_code = row['ts_code']
            ann_date = row['ann_date'].strftime('%Y-%m-%d')
            query = f"""
            SELECT trade_date, closes FROM date_1 WHERE date_1.trade_date > '{ann_date}' and date_1.ts_code = '{ts_code}' 
            order by trade_date asc limit 20
            """
            
            cursor.execute(query)
            result = cursor.fetchall()
            df1 = pd.DataFrame(result)
         
            try: 
                max_close = df1['closes'].max()
                min_close = df1['closes'].min()
                the_close = df1['closes'].iloc[-1]
                
                new_list.append({
                    'ts_code': ts_code,
                    'ann_date': ann_date,
                    'max_close': max_close,
                    'min_close': min_close,
                    'the_close': the_close,
                    'eps': row['eps'],
                    'total_revenue_ps': row['total_revenue_ps'],
                    'undist_profit_ps': row['undist_profit_ps'],
                    'gross_margin': row['gross_margin'],
                    'fcff': row['fcff'],
                    'fcfe': row['fcfe'],
                    'tangible_asset': row['tangible_asset'],
                    'bps': row['bps'],
                    'grossprofit_margin': row['grossprofit_margin'],
                    'npta': row['npta'],
                })
            except Exception as e:
                print(e)
            
        df2 = pd.DataFrame(new_list)
        df2.to_csv('fina_indicator.csv', index=False)
            
if __name__ == '__main__':
    mu = MySqlUtils()
    conn = mu.conn
    c1 = ClassIfcation()
    
    df = c1.get_fina_indicator(conn)
    c1.get_daily(conn, df)
